import torch 
import cv2
import os 
import numpy as np
from PIL import Image
from pytorch_grad_cam import GradCAM, ScoreCAM, GradCAMPlusPlus, AblationCAM, XGradCAM, EigenCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
from pytorch_grad_cam.utils.find_layers import find_layer_types_recursive
from torchvision import transforms
model_path = 'log/model/resnet50_0.pth' 
image_dir = '../data/dataset_resnet_single/0'
from tqdm import tqdm

transforms = transforms.Compose([
            transforms.CenterCrop((96, 96)),
            transforms.ToTensor(),
        ])
model = torch.load(model_path)
# model.eval()
for image_name in tqdm(os.listdir(image_dir)):
    image_path = os.path.join(image_dir,image_name)


    # print (model)
    # print (model.named_parameters())
    # for layer in model.named_parameters():
    #     print (layer[0])
    # for layer in find_layer_types_recursive(model, [torch.nn.Conv2d]):
    #     print (layer)
    target_layers = [find_layer_types_recursive(model, [torch.nn.Conv2d])[-1]]
    image = Image.open(image_path)
    image = transforms(image)
    input_tensor = image.unsqueeze(0)
    # print (np.max(input_tensor))
    cam = GradCAM(model=model, target_layers=target_layers, use_cuda=True)
    target_category = 0

    # You can also pass aug_smooth=True and eigen_smooth=True, to apply smoothing.
    grayscale_cam = cam(input_tensor=input_tensor, target_category=target_category)

    # In this example grayscale_cam has only one image in the batch:
    grayscale_cam = grayscale_cam[0, :]
    
    # print (np.array(image)[1,:,:].shape)
    image = cv2.cvtColor(np.array(image)[0,:,:], cv2.COLOR_GRAY2RGB)
    # print (image.shape)
    visualization = show_cam_on_image(image, grayscale_cam, use_rgb=True)

    Image.fromarray(visualization).save('grad_cam/'+image_name)

